import numpy as np

def shift(arr, num, fill_value=np.nan):
    arr = np.roll(arr, num)
    if num < 0:
        arr[num:] = fill_value
    elif num > 0:
        arr[:num] = fill_value
    return arr


def filter_pred(values_pred, scale=3):
    """通过比较相邻的异常点来确定异常点的范围。然后，根据指定的比例(scale)
       来过滤掉长度小于等于scale-1的异常点。

    Args:
        values_pred (np.ndarray): 原始预测
        scale (int, optional): 范围. Defaults to 3.

    Returns:
        ndarray: 修改后的预测值
    """
    predicted_anomalies_ = np.argwhere(values_pred == 1).ravel()
    predicted_anomalies_shift_forward = shift(predicted_anomalies_, 1, fill_value=predicted_anomalies_[0])
    predicted_anomalies_shift_backward = shift(predicted_anomalies_, -1, fill_value=predicted_anomalies_[-1])
    predicted_anomalies_start = np.argwhere(
        (predicted_anomalies_shift_forward - predicted_anomalies_) != -1
    ).ravel()
    predicted_anomalies_finish = np.argwhere(
        (predicted_anomalies_ - predicted_anomalies_shift_backward) != -1
    ).ravel()
    predicted_anomalies = np.hstack(
        [
            predicted_anomalies_[predicted_anomalies_start].reshape(-1, 1),
            predicted_anomalies_[predicted_anomalies_finish].reshape(-1, 1),
        ]
    )
    for a_range in predicted_anomalies:
        if a_range[1]-a_range[0]<=scale-1:
            values_pred[a_range[0]:a_range[1]+1] = 0
    return values_pred


def calc_point2point(predict, actual):
    """
    通过预测值和真实标签计算指标.

    Args:
        predict (np.ndarray): 预测值
        actual (np.ndarray): 真实标签
    """
    TP = np.sum(predict * actual)
    TN = np.sum((1 - predict) * (1 - actual))
    FP = np.sum(predict * (1 - actual))
    FN = np.sum((1 - predict) * actual)
    precision = TP / (TP + FP + 0.00001)
    recall = TP / (TP + FN + 0.00001)
    f1 = 2 * precision * recall / (precision + recall + 0.00001)
    return f1, precision, recall, TP, TN, FP, FN

def adjust_predicts(score, label,
                    threshold=None,
                    pred=None,
                    calc_latency=False,
                    is_filter=False):
    """
    根据给定的`score`、`threshold`（或给定的`pred`）和`label`计算调整后的预测标签。

    Args:
        score (np.ndarray): 异常分数
        label (np.ndarray): 实际标签
        threshold (float): 异常分数的阈值。如果分数低于阈值，则将该点标记为“异常”。
        pred (np.ndarray or None): 若不为None，则调整`pred`并忽略`score`和`threshold`，
        calc_latency (bool): 是否计算延迟

    Returns:
        np.ndarray: 预测标签
    """
    if len(score) != len(label):
        raise ValueError("score and label must have the same length")
    score = np.asarray(score)
    label = np.asarray(label)
    latency = 0
    if pred is None:
        predict = score > threshold
        if is_filter:
            predict = filter_pred(predict, scale=1)
    else:
        predict = pred
    actual = label > 0.1
    anomaly_state = False
    anomaly_count = 0
    for i in range(len(score)):
        if actual[i] and predict[i] and not anomaly_state:
            anomaly_state = True
            anomaly_count += 1
            for j in range(i, 0, -1):
                if not actual[j]:
                    break
                else:
                    if not predict[j]:
                        predict[j] = True
                        latency += 1
        elif not actual[i]:
            anomaly_state = False
        if anomaly_state:
            predict[i] = True
    if calc_latency:
        return predict, latency / (anomaly_count + 1e-4)
    else:
        return predict

def bf_search(label, score, verbose=True, is_filter=False, is_adjust = True):
    """寻找最佳的f1-score

    Args:
        label (np.ndarray): 测试集标签
        score (np.ndarray): 一场分数
        verbose (bool, optional): 打印额外信息. Defaults to True.
        is_filter (bool, optional): 过滤异常. Defaults to False.
        is_adjust (bool, optional): 是否使用点调整. Defaults to True.
    Returns:
        dict: 结果字典
    """
    start = 90
    search_range = [np.percentile(score, q) for q in np.arange(start, 100, 0.1)]
    m = {
        'f1-score': 0,
        'precision': 0,
        'recall': 0,
        'TP': 0,
        'TN': 0,
        'FP': 0,
        'FN': 0
    }
    for threshold in sorted(search_range)[::-1]:
        if is_adjust:
            pred = adjust_predicts(score, label, threshold, is_filter=is_filter)
        else:
            pred = score > threshold
        if len(np.unique(pred))==1:
            continue
        target = calc_point2point(pred, label)
        if target[0] >= m['f1-score']:
            m['f1-score'] = target[0]
            m['precision'] = target[1]
            m['recall'] = target[2]
            m['TP'] = target[3]
            m['TN'] = target[4]
            m['FP'] = target[5]
            m['FN'] = target[6]
            m['threshold'] = threshold
            if verbose:
                print("cur thr: ",target, m)
    return m